#!/bin/bash

echo "========================================"
echo "FLUX.1 训练和采样测试脚本"
echo "========================================"

# 设置环境变量
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export CUDA_LAUNCH_BLOCKING=1
export TOKENIZERS_PARALLELISM=false

# 基础路径定义
WORKSPACE="/workspace"
AI_DIR="$WORKSPACE/AI"
PYTHON_PATH="$AI_DIR/python/bin/python"
SCRIPT_PATH="$AI_DIR/ComfyUI/custom_nodes/comfyui_lora_train/sd_scripts/flux_train_network.py"

# 模型路径定义
MODELS_DIR="$AI_DIR/ComfyUI/models"
FLUX_MODEL="$MODELS_DIR/checkpoints/flux1-dev.safetensors"
CLIP_L_MODEL="$MODELS_DIR/clip/clip_l.safetensors"
T5XXL_MODEL="$MODELS_DIR/clip/t5xxl_fp16.safetensors"
AE_MODEL="$MODELS_DIR/vae/ae.safetensors"

# 数据集和输出路径
TRAIN_DATA_DIR="$WORKSPACE/training_data_all/girl"
OUTPUT_DIR="$WORKSPACE/training_data_all/output"
SAMPLE_PROMPTS="$WORKSPACE/training_data_all/sample_prompts.txt"

# 图片描述文件扩展名
CAPTION_EXTENSION=".txt"

# 训练参数配置
RESOLUTION="720,1280"

# 检查必要文件是否存在
echo "检查必要文件..."

if [ ! -f "$FLUX_MODEL" ]; then
    echo "错误: FLUX模型文件不存在: $FLUX_MODEL"
    exit 1
fi

if [ ! -f "$CLIP_L_MODEL" ]; then
    echo "错误: CLIP-L模型文件不存在: $CLIP_L_MODEL"
    exit 1
fi

if [ ! -f "$T5XXL_MODEL" ]; then
    echo "错误: T5XXL模型文件不存在: $T5XXL_MODEL"
    exit 1
fi

if [ ! -f "$AE_MODEL" ]; then
    echo "错误: AE模型文件不存在: $AE_MODEL"
    exit 1
fi

if [ ! -d "$TRAIN_DATA_DIR" ]; then
    echo "错误: 训练数据目录不存在: $TRAIN_DATA_DIR"
    exit 1
fi

if [ ! -f "$SAMPLE_PROMPTS" ]; then
    echo "错误: 采样提示词文件不存在: $SAMPLE_PROMPTS"
    exit 1
fi

# 创建输出目录
if [ ! -d "$OUTPUT_DIR" ]; then
    echo "创建输出目录: $OUTPUT_DIR"
    mkdir -p "$OUTPUT_DIR"
fi

echo
echo "========================================"
echo "开始FLUX.1训练和采样测试"
echo "========================================"
echo "模型: $FLUX_MODEL"
echo "训练数据: $TRAIN_DATA_DIR"
echo "输出目录: $OUTPUT_DIR"
echo "采样提示词: $SAMPLE_PROMPTS"
echo

# 验证依赖库
echo "验证依赖库..."
"$PYTHON_PATH" -c "from huggingface_hub import cached_download, hf_hub_download; print('✓ cached_download和hf_hub_download都可以正常导入')"

if [ $? -ne 0 ]; then
    echo "❌ 依赖库验证失败，请先安装兼容版本"
    echo "运行: pip install 'huggingface-hub==0.24.5' 'diffusers[torch]==0.25.0'"
    exit 1
fi

# 执行训练命令 (使用accelerate launch，符合README示例)
echo "执行训练命令..."

# 设置环境变量确保使用正确的Python
export PATH="$AI_DIR/python/bin:$PATH"

accelerate launch --mixed_precision bf16 --num_cpu_threads_per_process 1 "$SCRIPT_PATH" \
    --pretrained_model_name_or_path="$FLUX_MODEL" \
    --clip_l="$CLIP_L_MODEL" \
    --t5xxl="$T5XXL_MODEL" \
    --ae="$AE_MODEL" \
    --train_data_dir="$TRAIN_DATA_DIR" \
    --output_dir="$OUTPUT_DIR" \
    --output_name="flux_test_lora" \
    --network_module=networks.lora_flux \
    --network_dim=32 \
    --network_alpha=16 \
    --network_train_unet_only \
    --learning_rate=1e-4 \
    --max_train_epochs=1 \
    --train_batch_size=1 \
    --resolution="$RESOLUTION" \
    --optimizer_type=AdamW \
    --lr_scheduler=constant_with_warmup \
    --lr_warmup_steps=50 \
    --mixed_precision=bf16 \
    --save_precision=bf16 \
    --save_every_n_steps=500 \
    --save_model_as=safetensors \
    --gradient_checkpointing \
    --max_data_loader_n_workers=2 \
    --persistent_data_loader_workers \
    --seed=42 \
    --max_grad_norm=1.0 \
    --flip_aug \
    --fp8_base \
    --blocks_to_swap=16 \
    --disable_mmap_load_safetensors \
    --cache_latents \
    --cache_text_encoder_outputs_to_disk \
    --sample_every_n_steps=200 \
    --sample_prompts="$SAMPLE_PROMPTS" \
    --sample_sampler=ddim \
    --sample_at_first \
    --caption_extension="$CAPTION_EXTENSION" \
    --timestep_sampling=shift \
    --discrete_flow_shift=3.1582 \
    --model_prediction_type=raw \
    --guidance_scale=1.0

echo
echo "========================================"
echo "训练完成！"
echo "输出目录: $OUTPUT_DIR"
echo "========================================" 